/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 适配GE新接口头文件提取
 * Author: lilianlin
 * Create: 2020-05-14
 */

#ifndef OP_HCOM_COMM_H
#define OP_HCOM_COMM_H

#include "comm.h"
#include "graph/utils/attr_utils.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/node_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/graph_utils.h"
#include "graph/node.h"
#include "common/optimizer/graph_optimizer.h"
#include "common/optimizer/graph_optimizer_types.h"
#include "graph/compute_graph.h"
#include "graph/ge_tensor.h"
#include "external/graph/types.h"
#include "hccl_ex.h"
#include "hcom_op_utils.h"
#include "es_private.h"
#include "hcom_pub.h"


namespace hccl {
/* 原图shap类型配置 */
const std::string ORIGINAL_GRAPH_SHAPE_TYPE = "original_graph_shape_type";
constexpr std::int64_t ORIGINAL_GRAPH_KNOWNSHAPE_TYPE = 0;
constexpr std::int64_t ORIGINAL_GRAPH_UNKNOWNSHAPE_TYPE = 1;
constexpr u32 ALLTOALLV_RANK_MAX_NUM = 256; // 受notify数量限制，全连接组网alltoallv最多支持256p 分级alltoallv可以做到512
constexpr u32 ALLTOALLVC_RANK_MAX_NUM = 256; // 受notify数量限制，全连接组网alltoallvc最多支持256p 分级alltoallv可以做到512
constexpr u32 TENSOR_ALIGNMENT_512 = 512;
constexpr u32 TENSOR_ALIGNMENT_32 = 32;
constexpr u32 UNIQUE_TAG_MAX_LEN = 64;

struct AlltoAllVParamsInfo {
    u64 sendCounts[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 sendDispls[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 recvCounts[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 recvDispls[ALLTOALLV_RANK_MAX_NUM] = {0};
    HcclDataType sendType = HCCL_DATA_TYPE_RESERVED;
    HcclDataType recvType = HCCL_DATA_TYPE_RESERVED;
};

struct AlltoAllVCParamsInfo {
    u64 sendCountMatrix[ALLTOALLVC_RANK_MAX_NUM * ALLTOALLVC_RANK_MAX_NUM] = {0};
    HcclDataType sendType = HCCL_DATA_TYPE_RESERVED;
    HcclDataType recvType = HCCL_DATA_TYPE_RESERVED;
};

struct ReduceScatterVParamsInfo {
    u64 sendCounts[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 sendDispls[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 recvCounts[ALLTOALLV_RANK_MAX_NUM] = {0};
};

struct AllGatherVParamsInfo {
    u64 sendCount[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 recvDispls[ALLTOALLV_RANK_MAX_NUM] = {0};
    u64 recvCounts[ALLTOALLV_RANK_MAX_NUM] = {0};
};

struct EmbeddingServiceParams {
    s32 tableId{};
    s32 tag{};
    s32 insertFlag{};
    s32 maxNum{};
    s32 embeddingDim{};
    s32 flags{}; // 是否统计key_count
    s32 maxEmbeddingDim{}; // embeddingDim最大值
    s8 uniqueTag[UNIQUE_TAG_MAX_LEN + 1] = {0};
};

using HCCL_KERNEL_INFO_PRIVATE_DEF = struct hcclKernelInfoPrivateDef {
    u8 group[GROUP_NAME_MAX_LEN + 1] = {0}; // 1为结束符预留
    size_t nodeNameHash = 0;
    size_t tensorNum = 0;
    size_t privateDefSize = 0;
    u32 graphId = 0;
    u32 srcRank = 0;
    u32 destRank = 0;
    u32 selfRank = 0;
    u32 srTag = 0;
    u32 originalGraphShapeType = 0;
    int64_t comm = 0;
    HcclDataType dataType = HCCL_DATA_TYPE_RESERVED;
    bool needMapRank = false;
    bool isOfflineComp = false; // 是否是离线编译
    DevType devType = DevType::DEV_TYPE_COUNT; // 只有离线编译时需要，在loadtask的时候校验一致性
    EmbeddingServiceParams esInfo;
};

using HCCL_ALLTOALLV_KERNEL_INFO_PRIVATE_DEF =
    struct hcclAlltoallvKernelInfoPrivateDef : HCCL_KERNEL_INFO_PRIVATE_DEF {
    AlltoAllVParamsInfo paramsInfo = AlltoAllVParamsInfo();
    AlltoAllVCParamsInfo cparamsInfo = AlltoAllVCParamsInfo();

    hcclAlltoallvKernelInfoPrivateDef(HCCL_KERNEL_INFO_PRIVATE_DEF KernelInfoPrivateDef)
    {
        for (u32 i = 0; i < GROUP_NAME_MAX_LEN + 1; i++) {
            this->group[i] = KernelInfoPrivateDef.group[i];
        }
        nodeNameHash = KernelInfoPrivateDef.nodeNameHash;
        tensorNum = KernelInfoPrivateDef.tensorNum;
        privateDefSize = KernelInfoPrivateDef.privateDefSize;
        graphId = KernelInfoPrivateDef.graphId;
        srcRank = KernelInfoPrivateDef.srcRank;
        destRank = KernelInfoPrivateDef.destRank;
        selfRank = KernelInfoPrivateDef.selfRank;
        srTag = KernelInfoPrivateDef.srTag;
        originalGraphShapeType = KernelInfoPrivateDef.originalGraphShapeType;
        comm = KernelInfoPrivateDef.comm;
        dataType = KernelInfoPrivateDef.dataType;
        needMapRank = KernelInfoPrivateDef.needMapRank;
        esInfo = KernelInfoPrivateDef.esInfo;
    };
};

using HCCL_REDUCESCATTERV_KERNEL_INFO_PRIVATE_DEF =
    struct hcclReduceScatterVKernelInfoPrivateDef : HCCL_KERNEL_INFO_PRIVATE_DEF {
    ReduceScatterVParamsInfo paramsInfo = ReduceScatterVParamsInfo();

    hcclReduceScatterVKernelInfoPrivateDef(HCCL_KERNEL_INFO_PRIVATE_DEF KernelInfoPrivateDef)
    {
        for (u32 i = 0; i < GROUP_NAME_MAX_LEN + 1; i++) {
            this->group[i] = KernelInfoPrivateDef.group[i];
        }
        nodeNameHash = KernelInfoPrivateDef.nodeNameHash;
        tensorNum = KernelInfoPrivateDef.tensorNum;
        privateDefSize = KernelInfoPrivateDef.privateDefSize;
        graphId = KernelInfoPrivateDef.graphId;
        srcRank = KernelInfoPrivateDef.srcRank;
        destRank = KernelInfoPrivateDef.destRank;
        selfRank = KernelInfoPrivateDef.selfRank;
        srTag = KernelInfoPrivateDef.srTag;
        originalGraphShapeType = KernelInfoPrivateDef.originalGraphShapeType;
        comm = KernelInfoPrivateDef.comm;
        dataType = KernelInfoPrivateDef.dataType;
        needMapRank = KernelInfoPrivateDef.needMapRank;
        esInfo = KernelInfoPrivateDef.esInfo;
    };
};

using HCCL_ALLGATHERV_KERNEL_INFO_PRIVATE_DEF =
    struct hcclAllGatherVKernelInfoPrivateDef : HCCL_KERNEL_INFO_PRIVATE_DEF {
    AllGatherVParamsInfo paramsInfo = AllGatherVParamsInfo();

    hcclAllGatherVKernelInfoPrivateDef(HCCL_KERNEL_INFO_PRIVATE_DEF KernelInfoPrivateDef)
    {
        for (u32 i = 0; i < GROUP_NAME_MAX_LEN + 1; i++) {
            this->group[i] = KernelInfoPrivateDef.group[i];
        }
        nodeNameHash = KernelInfoPrivateDef.nodeNameHash;
        tensorNum = KernelInfoPrivateDef.tensorNum;
        privateDefSize = KernelInfoPrivateDef.privateDefSize;
        graphId = KernelInfoPrivateDef.graphId;
        srcRank = KernelInfoPrivateDef.srcRank;
        destRank = KernelInfoPrivateDef.destRank;
        selfRank = KernelInfoPrivateDef.selfRank;
        srTag = KernelInfoPrivateDef.srTag;
        originalGraphShapeType = KernelInfoPrivateDef.originalGraphShapeType;
        comm = KernelInfoPrivateDef.comm;
        dataType = KernelInfoPrivateDef.dataType;
        needMapRank = KernelInfoPrivateDef.needMapRank;
        esInfo = KernelInfoPrivateDef.esInfo;
    };
};

const std::vector<std::string>  HCOM_SUPPORTED_OP_TYPE = {
    HCCL_KERNEL_OP_TYPE_BROADCAST,
    HCCL_KERNEL_OP_TYPE_REDUCE,
    HCCL_KERNEL_OP_TYPE_ALLREDUCE,
    HCCL_KERNEL_OP_TYPE_ALLGATHER,
    HCCL_KERNEL_OP_TYPE_ALLGATHERV,
    HCCL_KERNEL_OP_TYPE_REDUCESCATTER,
    HCCL_KERNEL_OP_TYPE_REDUCESCATTERV,
    HCCL_KERNEL_OP_TYPE_SEND,
    HCCL_KERNEL_OP_TYPE_RECEIVE,
    HCCL_KERNEL_OP_TYPE_REMOTE_READ,
    HCCL_KERNEL_OP_TYPE_REMOTE_WRITE,
    HCCL_KERNEL_OP_TYPE_REMOTE_REF_READ,
    HCCL_KERNEL_OP_TYPE_REMOTE_SCATTER_WRITE,
    HCCL_KERNEL_OP_TYPE_ALLTOALLV,
    HCCL_KERNEL_OP_TYPE_ALLTOALLVC,
    HCCL_KERNEL_OP_TYPE_GATHER_ALLTOALLV,
    HCCL_KERNEL_OP_TYPE_ALLTOALL,
    HCCL_KERNEL_OP_TYPE_REMOTE_LOOKUP,
    HCCL_KERNEL_OP_TYPE_COLL_REMOTE_UPDATE,
    HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP,
    HCCL_KERNEL_OP_TYPE_GATHER,
    HCCL_KERNEL_OP_TYPE_COLL_REMOTE_UPDATE_PAIRED,
    HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_PAIRED,
    HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_UNIQUED_PAIRED
};

HcclResult GetHcclCommInitStatus(bool& initiated);
HcclResult GetDeviceType(const char *group, DevType &deviceType);
HcclResult GetRankSize(const char *group, u32 *rankSize);
HcclResult GetHccsLinkNum(const char *group, u32 &numHccsLink);
HcclResult GetHcclRankId(const char *group, u32 *rankId);
HcclResult UnloadHcclTask(const std::string group, const char *tag);
HcclResult SetWorkspaceResource(const std::string &tag, const char *group, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize);
HcclResult GetTopoDesc(const char *group, HcclTopoDescs *topoDescs, uint32_t topoSize);

HcclResult HcomInitialize();
HcclResult GetEsMaxRemoteOpNumPerStreamFromGe();
HcclResult GetOpDescIntAttr(const ge::OpDesc &op, const string &attr, s32 &output);
HcclResult GetOpDescStrAttr(const ge::OpDesc &op, const string &attr, string &output);
HcclResult HcclGetRemoteOperationParams(const ge::OpDesc &opdesc, HcomRemoteOperationParams *opParams);
HcclResult InitGroup();
bool IsOpESScene();
}
#endif
